#!/usr/bin/python
#coding=utf8
"""
# Author: Kellan Fan
# Created Time : Fri 22 Dec 2023 07:27:11 PM CST

# File Name: getswappid.py
# Description:

"""
import argparse
from pathlib import Path

proc_path = Path('/proc')
all_process_swap_usage = []


def get_total_swap_mb():
    total_swap = 0
    for line in open('/proc/swaps').readlines():
        if line.startswith('/pitrix'):
            total_swap += int(line.split()[2])
    return total_swap / 1024


def get_all_process_swap_usage(cmdline_max_length: int = 40):
    total_swap_mb = get_total_swap_mb()
    if total_swap_mb == 0:
        print("No swap found!")
        return []

    for pid_path in [path for path in proc_path.iterdir() if path.is_dir() and path.name.isdigit()]:
        pid = int(pid_path.name)
        cmdline = ' '.join(open(str(pid_path / 'cmdline')).read().split('\x00')).strip() if (pid_path / 'cmdline').exists() else ''
        cmdline = cmdline[:cmdline_max_length] + '...' if cmdline_max_length != 0 and len(cmdline) > cmdline_max_length else cmdline
        smaps = pid_path / 'smaps'
        swap_size = round(sum([int(line.split(':')[1].split()[0]) for line in open(str(smaps)).readlines() if 'Swap:' in line]) / 1024, 2) if smaps.exists() else 0
        if swap_size != 0:
            swap_percent = round(swap_size * 100 / total_swap_mb, 2)
            all_process_swap_usage.append((pid, swap_percent, swap_size, cmdline))

    all_process_swap_usage.sort(key=lambda x: x[1])
    return all_process_swap_usage


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-t', '--top', type=int, default=0, help='top n process swap status, defaults to 0(unlimited)', dest='top')
    parser.add_argument('-l', '--cmdline-max-length', type=int, default=0, help='max length of cmdline, defaults to 0(unlimited)', dest='cmdline_max_length')
    args = parser.parse_args()
    all_process_swap_usage = get_all_process_swap_usage(cmdline_max_length=args.cmdline_max_length)[-args.top:]
    for process_swap_usage in all_process_swap_usage:
        print('PID: %10s, SWAP(%%): %3.2f%%, SWAP(MB): %10.2f, CMDLINE: %-s' % process_swap_usage)
    print('Total Swap: %s MB' % round(get_total_swap_mb(), 2))
